import numpy as np
import sys
import os
from typing import Tuple, Optional


def parse_args() -> Tuple[str, int, int]:
    """解析命令行参数"""
    if len(sys.argv) != 4:
        raise ValueError("Usage: python nv12_to_rgb.py <nv12_file> <width> <height>")
    
    nv12_path = sys.argv[1]
    try:
        width = int(sys.argv[2])
        height = int(sys.argv[3])
    except ValueError:
        raise ValueError("Width and height must be integers.")
    
    if width <= 0 or height <= 0:
        raise ValueError("Width and height must be positive integers.")
    
    return nv12_path, width, height

def load_nv12_file(file_path: str, width: int, height: int) -> np.ndarray:
    """加载NV12文件为numpy数组"""
    # 计算预期文件大小 (YUV420/NV12格式为1.5字节/像素)
    expected_size = int(width * height * 1.5)
    
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"NV12 file not found: {file_path}")
    
    file_size = os.path.getsize(file_path)
    if file_size != expected_size:
        raise ValueError(f"Invalid file size. Expected {expected_size} bytes, got {file_size} bytes.")
    
    # 将二进制数据读取为uint8数组并reshape为(height*3/2, width)
    nv12_data = np.fromfile(file_path, dtype=np.uint8)
    return nv12_data.reshape((height * 3 // 2, width))

def nv12_to_rgb(nv12_data, width=None, height=None):
    """
    将NV12格式的数据转换为RGB格式（纯NumPy实现）
    
    参数:
        nv12_data: 二维数组，形状为 (height*3/2, width)
        width: 图像宽度（可选）
        height: 图像高度（可选）
    
    返回:
        RGB图像 (height × width × 3)
    """
    if width is None or height is None:
        total_height = nv12_data.shape[0]
        height = int(total_height * 2 / 3)
        width = nv12_data.shape[1]
    
    # 分离Y和UV分量
    y = nv12_data[:height, :]
    uv = nv12_data[height:, :].reshape((-1, width))
    
    # 分离U和V分量并归一化到[-128,127]
    u = uv[:, 0::2].astype(np.float32) - 128
    v = uv[:, 1::2].astype(np.float32) - 128
    
    # 上采样UV分量到Y的尺寸
    u = np.repeat(np.repeat(u, 2, axis=0), 2, axis=1)[:height, :width]
    v = np.repeat(np.repeat(v, 2, axis=0), 2, axis=1)[:height, :width]
    
    # YUV转RGB公式
    y = y.astype(np.float32)
    r = np.clip(y + 1.402 * v, 0, 255)
    g = np.clip(y - 0.344136 * u - 0.714136 * v, 0, 255)
    b = np.clip(y + 1.772 * u, 0, 255)
    
    return np.stack([r, g, b], axis=-1).astype(np.uint8)

def save_image(rgb: np.ndarray, output_path: str) -> None:
    """保存RGB图像为文件（使用OpenCV）"""
    import cv2
    # 将RGB转换为BGR格式（OpenCV默认使用BGR）
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, bgr)
    print(f"RGB image saved to: {output_path}")

def main() -> None:
    try:
        # 1. 解析命令行参数
        nv12_path, width, height = parse_args()
        
        # 2. 加载NV12文件
        print(f"Loading NV12 file: {nv12_path} [{width}x{height}]")
        nv12_data = load_nv12_file(nv12_path, width, height)
        
        # 3. 转换为RGB格式
        print("Converting NV12 to RGB...")
        rgb_image = nv12_to_rgb(nv12_data, width, height)
        
        # 4. 保存结果
        output_path = os.path.splitext(nv12_path)[0] + "_converted.jpg"
        save_image(rgb_image, output_path)
    
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()